package cn.seecoder.ai.serviceImpl;

import cn.seecoder.ai.algorithm.ml.BaseAlgorithm;
import cn.seecoder.ai.algorithm.ml.classification.*;
import cn.seecoder.ai.algorithm.ml.clustering.KMeansClustering;
import cn.seecoder.ai.algorithm.ml.regression.DecisionTreeRegression;
import cn.seecoder.ai.algorithm.ml.regression.RandomForestRegression;
import cn.seecoder.ai.dao.ModelRepository;
import cn.seecoder.ai.enums.MachineLearningAlgorithmEnum;
import cn.seecoder.ai.enums.MachineLearningAlgorithmTypeEnum;
import cn.seecoder.ai.enums.TrainParamEnum;
import cn.seecoder.ai.enums.TrainingStateEnum;
import cn.seecoder.ai.exception.AIExternalException;
import cn.seecoder.ai.model.bo.TrainParamsBO;
import cn.seecoder.ai.model.po.Model;
import cn.seecoder.ai.model.vo.ConfigVO;
import cn.seecoder.ai.model.vo.ModelVO;
import cn.seecoder.ai.service.ConfigService;
import cn.seecoder.ai.service.FeatureEngineeringService;
import cn.seecoder.ai.service.ModelService;
import cn.seecoder.ai.utils.HdfsHelper;
import cn.seecoder.ai.utils.POJOConverter;
import lombok.extern.slf4j.Slf4j;
import org.apache.spark.ml.PipelineModel;
import org.apache.spark.sql.SparkSession;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import scala.Tuple2;

import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

import static cn.seecoder.ai.utils.MLAlgorithmHelper.getMachineLearningAlgorithmTypeEnum;


/**
 * @author fanyanpeng
 * @date 2023/4/17 2:05
 */
@Slf4j
@Service
public class ModelServiceImpl implements ModelService  {




    @Autowired
    ConfigService configService;

    @Autowired
    ModelRepository modelRepository;

    @Autowired
    FeatureEngineeringService featureEngineeringService;

    @Autowired
    SparkSession sparkSession;


    @Autowired
    HdfsHelper hdfsHelper;


    /**
     * 创建一个模型
     * @author   fanyanpeng
     * @date 2023/4/17 4:57
     * @param modelVO
     * @return cn.seecoder.ai.model.vo.ModelVO
     */
    @Override
    public ModelVO createModel(ModelVO modelVO) {

        if(modelVO.getTrainParams()==null){
            modelVO.setTrainParams(new HashMap<>());
        }
        ensureAlgorithmMatchesTrainParams(modelVO.getMachineLearningAlgorithm(),modelVO.getTrainParams());
        Model model = toPO(modelVO);
        model.setCreatedTime(new Date());
        model.setTrainingState(TrainingStateEnum.WAITING);
        model = modelRepository.save(model);
        return toVO(model);
    }

    /**
     * 获得某用户的所有模型
     * @param userId
     * @return
     */
    @Override
    public List<ModelVO> getModelVOList(Integer userId) {
        List<Model> modelList = modelRepository.findByUserIdOrderByCreatedTimeDesc(userId);
        return modelList.stream().map(this::toVO).collect(Collectors.toList());
    }

    /**
     * 某个配置下，某个用户的所有模型
     * @param userId
     * @param configId
     * @return
     */
    @Override
    public List<ModelVO> getModelVOList(Integer userId, Integer configId) {
        List<Model> modelList = modelRepository.findByUserIdAndConfigId(userId,configId);
        return modelList.stream().map(this::toVO).collect(Collectors.toList());
    }

    /**
     * 获取单个模型
     * @param userId
     * @param modelId
     * @return
     */
    @Override
    public ModelVO getModelVO(Integer userId, Integer modelId) {
        Optional<Model>  modelOptional= modelRepository.findById(modelId);
        if(!modelOptional.isPresent()){
            throw  AIExternalException.modelNotFoundById(modelId);
        }
        Model model = modelOptional.get();
        return toVO(model);
    }





    public ModelVO renewModel(Integer userId, Integer modelId){
        Model model = getModel(modelId);
        model.setTrainingState(TrainingStateEnum.WAITING);
        model.setTrainingStartTime(null);
        model.setTrainingFinishTime(null);
        model.setFeaturesIndexerUri(null);
        model.setLabelIndexerUri(null);
        model.setModelUri(null);
        model.setTrainResults(null);
        modelRepository.save(model);
        return getModelVO(userId, modelId);
    }



    private String feature(Integer modelId, Integer configId) throws IOException {
        Model model = getModel(modelId);
        MachineLearningAlgorithmEnum machineLearningAlgorithmEnum = model.getMachineLearningAlgorithm();
        ConfigVO configVO = configService.getConfig(0,configId);
        // 获得训练文件id
        Integer fileInfoId = configVO.getFileInfoId();

        String libSvmUri = null;

        // 如果填写了测试文件id
        if(configVO.getTestFileInfoId()!=null){
            Integer testFileInfoId = configVO.getTestFileInfoId();
            // 生成训练文件的libSvm格式地址
            libSvmUri= featureEngineeringService.processTrainAndTestFile(
                    configId,
                    fileInfoId,
                    testFileInfoId,
                    hdfsHelper.getUriLibSvm(modelId,fileInfoId),
                    hdfsHelper.getUriFeatureStringIndexerPipeline(modelId),
                    hdfsHelper.getUriLabelStringIndexer(modelId),
                    machineLearningAlgorithmEnum);
            return libSvmUri;
        }

        // 如果没有填写测试文件id
        libSvmUri = featureEngineeringService.processSingleFile(
                configId,
                fileInfoId,
                hdfsHelper.getUriLibSvm(modelId,fileInfoId),
                hdfsHelper.getUriFeatureStringIndexerPipeline(modelId),
                hdfsHelper.getUriLabelStringIndexer(modelId),
                machineLearningAlgorithmEnum,
                true);
        return libSvmUri;
    }

    public ModelVO syncTrainModel(Integer userId, Integer modelId){
        // 更新机器学习模型状态，进入训练
        renewModel(userId, modelId);
        Model model = getModel(modelId);
        try {
            // 设置对象状态为正在训练，设置训练开始时间
            model.setTrainingState(TrainingStateEnum.ON_GOING);
            model.setTrainingStartTime(new Date());
            modelRepository.save(model);

            // 机器学习算法
            MachineLearningAlgorithmEnum machineLearningAlgorithmEnum = model.getMachineLearningAlgorithm();
            // 获得配置
            Integer configId = model.getConfigId();

            // 生成训练文件的libSvm格式地址
            String libSvmUri = feature(modelId, configId);

            // 更新map中的trainSetOccupy，训练集占比
            model = getModel(modelId);

            PipelineModel  pipelineModel = null;
            Map<String,Object>  trainResults = null;
            /* 使用机器学习算法进行训练 */
            Tuple2<PipelineModel,Map<String,Object>> result = train(model.getTrainParams(), libSvmUri, machineLearningAlgorithmEnum);

            pipelineModel = result._1;  // 结果一：流水线模型
            trainResults = result._2;   // 结果二：Map<String,Object>格式的训练结果

            // 保存模型到HDFS
            String modelUri= savePipelineModelToHdfs(pipelineModel,modelId,userId);

            // 设置对象状态为训练完成，设置训练结束时间
            model.setModelUri(modelUri);
            model.setTrainResults(trainResults);
            model.setTrainingState(TrainingStateEnum.SUCCESS);
            modelRepository.save(model);
            log.info(String.format("%s model train finished, result is as follows: %s",model.getMachineLearningAlgorithm().getAlgorithmName(),model.getTrainResults().toString()));

        }catch (Exception e){   //训练出错
            e.printStackTrace();
            // 设置对象状态为训练失败，设置训练结束时间
            model.setTrainingState(TrainingStateEnum.FAIL);
            modelRepository.save(model);
        }
        return getModelVO(userId, modelId);
    }

    @Async
    public void asyncTrainModel(Integer userId, Integer modelId){
        syncTrainModel(userId, modelId);
    }

    /**
     *
     * @author   fanyanpeng
     * @date 2023/4/22 22:32
     * @param trainParams 训练参数
     * @param libSvmUri 训练文件的libsvm格式地址
     * @param machineLearningAlgorithmEnum 机器学习算法
     * @return scala.Tuple2<org.apache.spark.ml.PipelineModel,java.util.Map<java.lang.String,java.lang.Object>>
     */
    private Tuple2<PipelineModel, Map<String, Object>> train(Map<TrainParamEnum,Object> trainParams,
                                                             String libSvmUri,
                                                             MachineLearningAlgorithmEnum machineLearningAlgorithmEnum) {

        BaseAlgorithm baseAlgorithm;
        switch (machineLearningAlgorithmEnum){
            case DECISION_TREE_CLASSIFICATION:{
                baseAlgorithm = new DecisionTreeClassification();
                break;
            }
            case GRADIENT_BOOSTED_TREE_CLASSIFICATION:{
                baseAlgorithm = new GradientBoostedTreeClassification();
                break;
            }
            case LOGISTIC_REGRESSION_CLASSIFICATION:{
                baseAlgorithm = new LogisticRegressionClassification();
                break;
            }
            case MULTILAYER_PERCEPTRON_CLASSIFICATION:{
                baseAlgorithm = new MultilayerPerceptronClassification();
                break;
            }
            case NAIVE_BAYES_CLASSIFICATION:{
                baseAlgorithm = new NaiveBayesClassification();
                break;
            }
            case RANDOM_FOREST_CLASSIFICATION:{
                baseAlgorithm = new RandomForestClassification();
                break;
            }

            case DECISION_TREE_REGRESSION:{
                baseAlgorithm = new DecisionTreeRegression();
                break;
            }

            case RANDOM_FOREST_REGRESSION:{
                baseAlgorithm = new RandomForestRegression();
                break;
            }

            case K_MEANS_CLUSTERING:{
                baseAlgorithm = new KMeansClustering();
                break;
            }
            case XGBOOST_CLASSIFICATION:{
                baseAlgorithm = new XGBoostClassification();
                break;
            }
            default:{
                throw AIExternalException.machineLearningAlgorithmNotMapped();
            }
        }

        return baseAlgorithm.train(sparkSession,libSvmUri,new TrainParamsBO(trainParams));
    }




    /**
     * 返回pipelineModel保存的地址。
     * @author   fanyanpeng
     * @date 2023/4/19 3:00
     * @param pipelineModel
     * @param modelId
     * @param userId
     * @return java.lang.String
     */
    private String savePipelineModelToHdfs(PipelineModel pipelineModel,Integer modelId,Integer userId){
        String modelUri = hdfsHelper.getUriPipelineModel(userId, modelId);
        try{
            pipelineModel.write().overwrite().save(modelUri);
        }catch (IOException e ){
            e.printStackTrace();
            throw AIExternalException.pipeLineSaveToHdfsFailed();
        }
        return modelUri;
    }



    private void addAdditionalName(ModelVO modelVO){
        try {
            ConfigVO configVO = configService.getConfig(0, modelVO.getConfigId());
            modelVO.setConfigName(configVO.getConfigName());
            modelVO.setFileInfoName(configVO.getFileInfoName());
        }catch (Exception e){
            log.warn(e.getMessage());
        }

    }


    private ModelVO toVO(Model model){
        ModelVO modelVO = (ModelVO) POJOConverter.convert(model);
        addAdditionalName(modelVO);
        return modelVO;
    }

    private Model toPO(ModelVO modelVO){
        return (Model) POJOConverter.convert(modelVO);
    }


    /**
     * 确保参数都在范围之内
     * @author   fanyanpeng
     * @date 2023/4/17 4:47
     * @param machineLearningAlgorithmEnum
     * @param trainParams
     * @return java.lang.Boolean
     */
    private Boolean ensureAlgorithmMatchesTrainParams(MachineLearningAlgorithmEnum machineLearningAlgorithmEnum,
                                                      Map<TrainParamEnum,Object> trainParams){
        TrainParamEnum[] trainParamEnumNeeded = machineLearningAlgorithmEnum.getTrainParamEnums();

        for(TrainParamEnum trainParamEnum: trainParamEnumNeeded){
            Object value = trainParams.get(trainParamEnum);
            if(value == null){
                continue;
            }
            // 判断范围
            trainParamEnum.validInRange(value);
        }
        return true;
    }

    private Model getModel(Integer modelId){
        return modelRepository.findById(modelId).get();
    }


}


